'****************************************
' Run an arbitrary number of FTP clients
'****************************************

Dim retCode : retCode = 1

If WScript.Arguments.Count < 2 Then
	WScript.Echo "Usage: runftp <digest|session> <ftp_server[</path>]> [<num_sessions> <files_per_session> [<timeout_sec>]]"
	WScript.Quit retCode
End If

Dim objShell : Set objShell = WScript.CreateObject("WScript.Shell")

'--------------- check if fciv.exe is present ---------------
Dim fFcivExists

fFcivExists = objShell.Run("cmd.exe /C fciv.exe -h", 0, True)

If fFcivExists <> 0 Then
	WScript.Echo "Couldn't run ""File checksum integrity verifier"" (fciv.exe)"
	WScript.Echo "please download and/or install `windows-kb841290-x86-enu.exe`"
	WScript.Echo "<http://download.microsoft.com/download/c/f/4/cf454ae0-a4bb-4123-8333-a1b6737712f7/windows-kb841290-x86-enu.exe>"
	WScript.Quit retCode
End If

Const ForReading = 1
Const ForWriting = 2
Const SecondsMin = 0.01
' name of file with checksums 
Dim fileList : fileList = "$file_list$"
' name of file with ftp commands
Dim ftpCommands : ftpCommands = "$ftp_commands$"
' name of temporary folder
Dim tmpDir : tmpDir = "$tmp_folder$"
' name of session folder 
Dim sesnDir : sesnDir = "$tmp_session" ' will be trailed with session number ("XX$")

Dim objFSO, objFtpCmds, objInteg, objResults, strEntry, objTmpDir

Set objFSO = CreateObject("Scripting.FileSystemObject")

Dim schemePos
Dim pathPos
Dim hostPart
Dim userPos : userPos = 0
Dim passwdPos : passwdPos = 0
Dim ftpHost : ftpHost = ""
Dim ftpPath : ftpPath = ""
Dim userLogin : userLogin = ""
Dim userPasswd : userPasswd = ""
Dim paraAnon : paraAnon = "-A "
Dim fDirEntry
Dim strParts
Dim fileNames : fileNames = Array()
Dim nameIndex, auxIndex
Dim strPart
Dim objCheck
Dim objFileHashes
Dim sesnInst : sesnInst = Array()
Dim dirName
Dim regBytesSec : Set regBytesSec = New RegExp
Dim regMatches
Dim currRate, maxRate, minRate, avgSesn, numRates
Dim avgRates : avgRates = Array()
Dim valTimeout : valTimeout = 600 ' as seconds
Dim timeCheck
Dim matchBytes
Dim regSeconds : Set regSeconds = New RegExp
Dim numFiles

'--------------- parse FTP server arg ---------------
' if someone accidentally typed scheme prefix or just copied url
schemePos = InStr(WScript.Arguments(1), "ftp://")

If schemePos > 0 Then
	schemePos = schemePos + 6
Else
	schemePos = 1
End If

pathPos = InStr(schemePos, WScript.Arguments(1), "/")

If pathPos > 0 Then
	ftpPath = Mid(WScript.Arguments(1), pathPos + 1)
	hostPart = Mid(WScript.Arguments(1), schemePos, pathPos - schemePos)
Else
	hostPart = Mid(WScript.Arguments(1), schemePos)
End If

' check for user:password
userPos = InStr(hostPart, "@")

If userPos > 0 Then
	ftpHost = Mid(hostPart, userPos + 1)
	userLogin = Left(hostPart, userPos - 1)

	passwdPos = InStr(userLogin, ":")

	If passwdPos > 0 Then
		userPasswd = Mid(userLogin, passwdPos + 1)
		userLogin = Left(hostPart, passwdPos - 1)
	Else
		userPasswd = ""
	End If
	paraAnon = ""
Else
	ftpHost = hostPart
End If

' init random seed
Randomize

'--------------- decide on command ---------------
If StrComp(WScript.Arguments(0), "digest") = 0 Then

	WScript.Echo "*** calculating file checksums ***"

	' get file list
	Set objFtpCmds = objFSO.OpenTextFile(ftpCommands, ForWriting, True)
	If userPos > 0 Then
		objFtpCmds.Write userLogin & vbCrLf
		If passwdPos > 0 Then
			objFtpCmds.Write userPasswd & vbCrLf
		End If
	End If
	If Len(ftpPath) > 0 Then
		objFtpCmds.Write "cd " & ftpPath & vbCrLf
	End If
	objFtpCmds.Write "dir" & vbCrLf
	objFtpCmds.Write "bye" & vbCrLf
	objFtpCmds.Close
	Set objFtpCmds = Nothing

	objResults = Split(RunFTPQuery, vbCrLf)

	fDirEntry = 0
	For Each strEntry In objResults
		If StrComp(Left(strEntry, 4), "150 ") = 0 Then
			fDirEntry = 1
		ElseIf StrComp(Left(strEntry, 4), "125 ") = 0 Then
			fDirEntry = 1
		ElseIf StrComp(Left(strEntry, 4), "226 ") = 0 Then
			Exit For
		ElseIf fDirEntry = 1 Then
			strParts = Split(strEntry, " ")
			' trim trainling CR
			strPart = strParts(UBound(strParts))
			strPart = Left(StrPart, Len(StrPart) - 1)
			nameIndex = UBound(fileNames)
			ReDim Preserve fileNames(nameIndex + 1)
			fileNames(nameIndex + 1) = strPart
		End If
	Next

	If fDirEntry = 0 Then
		WScript.Echo "error getting file list"
		WScript.Quit retCode
	End If

	' create temporary folder
	If objFSO.FolderExists(tmpDir) Then
		objFSO.DeleteFolder(tmpDir)
	End If

	objTmpDir = objFSO.CreateFolder(tmpDir)

	' get all the available files
	Set objFtpCmds = objFSO.OpenTextFile(ftpCommands, ForWriting, True)
	If userPos > 0 Then
		objFtpCmds.Write userLogin & vbCrLf
		If passwdPos > 0 Then
			objFtpCmds.Write userPasswd & vbCrLf
		End If
	End If
	If Len(ftpPath) > 0 Then
		objFtpCmds.Write "cd " & ftpPath & vbCrLf
	End If
	objFtpCmds.Write "lcd " & tmpDir & vbCrLf
	objFtpCmds.Write "binary" & vbCrLf
	For Each strEntry In fileNames
		objFtpCmds.Write "get " & strEntry  & vbCrLf
	Next
	objFtpCmds.Write "bye" & vbCrLf
	objFtpCmds.Close
	Set objFtpCmds = Nothing

	' UNDONE: check every download for success
	objResults = RunFTPQuery

	' calculate files checksum
	' UNDONE: use fciv.exe xml db format
	Set objCheck = objShell.Exec("fciv -wp " & tmpDir) ' exclude basedir from filename
	objResults = Split(objCheck.StdOut.ReadAll, vbCrLf)

	' write down filenames along with hashes
	Set objInteg = objFSO.OpenTextFile(fileList, ForWriting, True)

	nameIndex = 0
	For Each strEntry In objResults
		If Len(strEntry) > 0 Then
			If StrComp(Left(strEntry, 2), "//") <> 0 Then
				strParts = Split(strEntry, " ")
				objInteg.Write strParts(1) & " " & strParts(0) & vbCrLf
				nameIndex = nameIndex + 1
			End If
		End If
	Next

	objInteg.Close

	' delete temporary folder
	If objFSO.FolderExists(tmpDir) Then
		objFSO.DeleteFolder(tmpDir)
	End If

	Set objTmpDir = Nothing

	' UNDONE: delete ftp commands file

	If nameIndex = 0 Then
		WScript.Echo "file list is empty"
		WScript.Quit retCode
	End If

	retCode = 0 ' success
	WScript.Echo "done"

ElseIf StrComp(WScript.Arguments(0), "session") = 0 Then

	WScript.Echo "*** running test sessions ***"

	' check arguments
	If WScript.Arguments.Count < 4 Then
		WScript.Echo "session parameters are missed"
		WScript.Quit retCode
	End If

	If WScript.Arguments(2) < 1 Or WScript.Arguments(3) < 1 Then
		WScript.Echo "session parameter is invalid"
		WScript.Quit retCode
	End If

	If WScript.Arguments.Count = 5 Then
		valTimeout = WScript.Arguments(4)
	End If

	If valTimeout < 1 Then
		WScript.Echo "timeout is invalid"
		WScript.Quit retCode
	End If

	WScript.Echo "with timeout = " & valTimeout & " sec."

	' load file list
	Set objFileHashes = CreateObject("Scripting.Dictionary")

	Set objInteg = objFSO.OpenTextFile(fileList, ForReading, True)

	Do While objInteg.AtEndOfStream <> True
		strEntry = objInteg.ReadLine
		strParts = Split(strEntry, " ")
		objFileHashes.Add strParts(0), strParts(1)
		nameIndex = UBound(fileNames)
		ReDim Preserve fileNames(nameIndex + 1)
		fileNames(nameIndex + 1) = strParts(0)
	Loop

	objInteg.Close

	' prepare sessions
	For nameIndex = 0 To WScript.Arguments(2) - 1
		dirName = sesnDir & nameIndex & "$"

		If objFSO.FolderExists(dirName) Then
			objFSO.DeleteFolder(dirName)
		End If
		objFSO.CreateFolder(dirName)
		
		Set objFtpCmds = objFSO.OpenTextFile(dirName & "\" & ftpCommands, ForWriting, True)
		If userPos > 0 Then
			objFtpCmds.Write userLogin & vbCrLf
			If passwdPos > 0 Then
				objFtpCmds.Write userPasswd & vbCrLf
			End If
		End If
		If Len(ftpPath) > 0 Then
			objFtpCmds.Write "cd " & ftpPath & vbCrLf
		End If
		objFtpCmds.Write "lcd " & dirName & vbCrLf
		objFtpCmds.Write "binary" & vbCrLf
		For auxIndex = 0 To WScript.Arguments(3) - 1
			objFtpCmds.Write "get " & fileNames(Round(Rnd * UBound(fileNames))) & vbCrLf
		Next		
		objFtpCmds.Write "bye" & vbCrLf
		objFtpCmds.Close
		Set objFtpCmds = Nothing
	Next

	' fire sessions
	For nameIndex = 0 To WScript.Arguments(2) - 1
		dirName = sesnDir & nameIndex & "$"
		auxIndex = UBound(sesnInst)
		ReDim Preserve sesnInst(auxIndex + 1)
		Set sesnInst(auxIndex + 1) = objShell.Exec("ftp -v " & paraAnon & "-s:" & dirName & "\" & ftpCommands & " " & ftpHost)
	Next

	regBytesSec.Pattern = "([\d]+)[^\d]+([\d]+[,\.][\d]+)?"
	regBytesSec.Global = False
	regBytesSec.Multiline = False

	regSeconds.Pattern = "^([\d]+[,\.][\d]{2})"
	regSeconds.Global = False
	regSeconds.Multiline = False

	maxRate = 0
	minRate = 999999999 ' FIXME: silly, rework
	Redim avgRates(WScript.Arguments(2))
	timeCheck = DateAdd("s", valTimeout, Now)
	numFiles = 0

	' wait for sessions to complete and gather stats
	auxIndex = WScript.Arguments(2)
	While auxIndex > 0 
		For nameIndex = 0 To WScript.Arguments(2) - 1
			If Not sesnInst(nameIndex) Is Nothing Then
				If sesnInst(nameIndex).Status <> 0 Then
					auxIndex = auxIndex - 1
					objResults = Split(sesnInst(nameIndex).StdOut.ReadAll, vbCrLf)
					numRates = 0
					avgSesn = 0
					fDirEntry = 0
					For Each strEntry In objResults
						strPart = Left(strEntry, 4)
						If StrComp(strPart, "421 ") = 0 Then
							WScript.Echo "error   : connection closed (421)"
							Exit For
						ElseIf StrComp(strPart, "425 ") = 0 Then
							WScript.Echo "error   : can't open data connection (425)"
							Exit For
						ElseIf StrComp(strPart, "426 ") = 0 Then
							WScript.Echo "error   : transfer aborted (426)"
							Exit For
						ElseIf StrComp(strPart, "226 ") = 0 Then
							fDirEntry = 1
						ElseIf fDirEntry = 1 Then
							If regBytesSec.Test(strEntry) Then
								Set regMatches = regBytesSec.Execute(strEntry)
								' count as Kbytes
								If Len(regMatches(0).SubMatches(1)) = 0 Then
									matchBytes = regMatches(0).SubMatches(0)
									fDirEntry = 2
								ElseIf CDbl(regMatches(0).SubMatches(1)) > SecondsMin Then
									currRate =  (regMatches(0).SubMatches(0) / CDbl(regMatches(0).SubMatches(1))) / 1000.0
									Set regMatches = Nothing
									fDirEntry = 3
								End If
							Else
								fDirEntry = 0
							End If
						ElseIf fDirEntry = 2 Then
							' Windows 7/Server 2008 ftp fucked-up stdout
							If regSeconds.Test(strEntry) Then
								Set regMatches = regSeconds.Execute(strEntry)
								If CDbl(regMatches(0).SubMatches(0)) > SecondsMin Then
									currRate =  (matchBytes / CDbl(regMatches(0).SubMatches(0))) / 1000.0
									fDirEntry = 3
								Else
									fDirEntry = 0
								End If
							Else
								fDirEntry = 0
							End If
						Else
							fDirEntry = 0
						End If

						If fDirEntry = 3 Then
							If currRate > maxRate Then
								maxRate = currRate
							End If

							If currRate < minRate Then
								minRate = currRate
							End If

							avgSesn = avgSesn + currRate
							numRates = numRates + 1

							fDirEntry = 0
						End If

					Next
					' close session
					Set sesnInst(nameIndex) = Nothing
					
					If numRates > 0 Then
						avgRates(nameIndex) = avgSesn / numRates
					Else
						avgRates(nameIndex) = 0
					End If

				ElseIf DateDiff("s", Now, timeCheck) < 0 Then
					' terminate session
					sesnInst(nameIndex).Terminate
					auxIndex = auxIndex - 1
					WScript.Echo "error   : terminating session by timeout"
					Set sesnInst(nameIndex) = Nothing

				Else
					WScript.Sleep 100
				End If

				If sesnInst(nameIndex) Is Nothing Then ' session finished
					' check downloads for integrity
					Set objCheck = objShell.Exec("fciv -wp " & sesnDir & nameIndex & "$")
					objResults = Split(objCheck.StdOut.ReadAll, vbCrLf)

					For Each strEntry In objResults
						If Len(strEntry) > 0 Then
							If StrComp(Left(strEntry, 2), "//") <> 0 Then
								strParts = Split(strEntry, " ")
								If objFileHashes.Exists(strParts(1)) Then
									If StrComp(strParts(0), objFileHashes.Item(strParts(1))) <> 0 Then
										WScript.Echo "error   : checksum failed for " & strParts(1)
									End If
									numFiles = numFiles + 1 ' number of files checked
								End If
							End If
						End If
					Next

					Set objCheck = Nothing
				End If
			End If ' session object is not null
		Next
	WEnd

	currRate = 0
	For nameIndex = 0 To WScript.Arguments(2) - 1
		currRate = currRate + avgRates(nameIndex)
	Next

	WScript.Echo "max rate: " & Round(maxRate, 3) & " Kbytes/sec"
	WScript.Echo "min rate: " & Round(minRate, 3) & " Kbytes/sec"
	WScript.Echo "average : " & Round(currRate / WScript.Arguments(2), 3) & " Kbytes/sec"

	WScript.Echo "files downloaded: " & numFiles

	' clean sessions
	For nameIndex = 0 To WScript.Arguments(2)
		dirName = sesnDir & nameIndex & "$"
		If objFSO.FolderExists(dirName) Then
			objFSO.DeleteFolder(dirName)
		End If
	Next

	Set objFileHashes = Nothing
Else
	WScript.Echo "Unknown command - " & WScript.Arguments(0)
End If

'--------------- cleanup ---------------
Set objResults = Nothing
Set objInteg = Nothing
Set objFtpCmds = Nothing
Set objFSO = Nothing
Set objShell = Nothing

WScript.Quit retCode

'--------------- RunFTPQuery ---------------
Function RunFTPQuery

	Dim objFTP : Set objFTP = objShell.Exec("ftp -v " & paraAnon & "-s:" & ftpCommands & " " & ftpHost)
	' UNDONE: wait on completion, may take some time with name resolving and such
	RunFTPQuery = objFTP.StdOut.ReadAll

	Set objFTP = Nothing

End Function

' change to temporary dir
'runDir = objShell.CurrentDirectory
'objShell.CurrentDirectory = runDir & "\" & tmpDir

'	For nameIndex = 0 To UBound(fileNames)
'		WScript.Echo fileNames(nameIndex) & " = " & objFileHashes.Item(fileNames(nameIndex))
'	Next

'WScript.Timeout = 1000
